import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import warnings
import datetime
import matplotlib.dates as mdates
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import FuncFormatter
import matplotlib.font_manager as fm
import platform
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import tushare as ts

# 设置中文字体
system = platform.system()
if system == "Windows":
    font_path = "C:/Windows/Fonts/SimHei.ttf"
elif system == "Darwin":  # macOS
    font_path = "/System/Library/Fonts/PingFang.ttc"
else:  # Linux
    font_path = "/usr/share/fonts/wqy-zenhei/wqy-zenhei.ttc"

prop = fm.FontProperties(fname=font_path)
plt.rcParams["font.family"] = prop.get_name()
plt.rcParams['axes.unicode_minus'] = False

# 初始化tushare接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


# 获取股票列表
def get_stock_list():
    # 获取所有A股股票
    stock_list = pro.stock_basic(exchange='', list_status='L',
                                 fields='ts_code,symbol,name,area,industry,list_date')
    # 筛选出一些有代表性的股票
    selected_stocks = [
        '000001.SZ', '600519.SH', '000651.SZ', '601318.SH', '600276.SH',
        '601166.SH', '000333.SZ', '600036.SH', '601888.SH', '600900.SH'
    ]
    return stock_list[stock_list['ts_code'].isin(selected_stocks)]


# 获取股票历史数据
def get_stock_data(stock_list, start_date='20200101', end_date='20230101'):
    all_data = []
    for ts_code in stock_list['ts_code']:
        try:
            # 获取日线数据
            df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
            # 获取复权因子用于计算真实价格
            adj_factor = pro.adj_factor(ts_code=ts_code, trade_date='')
            # 合并数据
            df = pd.merge(df, adj_factor, on=['ts_code', 'trade_date'], how='left')
            # 计算复权价格
            df = df.sort_values('trade_date')
            df['adj_factor'] = df['adj_factor'].fillna(method='bfill')
            df['adj_factor'] = df['adj_factor'] / df['adj_factor'].iloc[-1]
            df['adj_close'] = df['close'] * df['adj_factor']
            df['adj_open'] = df['open'] * df['adj_factor']
            df['adj_high'] = df['high'] * df['adj_factor']
            df['adj_low'] = df['low'] * df['adj_factor']
            all_data.append(df)
            print(f"成功获取 {ts_code} 的数据")
        except Exception as e:
            print(f"获取 {ts_code} 的数据时出错: {e}")
    return pd.concat(all_data, ignore_index=True)


# 获取财务指标数据
def get_financial_data(stock_list, start_date='20200101', end_date='20230101'):
    financial_data = []
    for ts_code in stock_list['ts_code']:
        try:
            # 获取PE和PB数据
            df = pro.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date,
                                 fields='ts_code,trade_date,pe,pb,ps,dv_ratio')
            financial_data.append(df)
        except Exception as e:
            print(f"获取 {ts_code} 的财务数据时出错: {e}")
    return pd.concat(financial_data, ignore_index=True)


# 数据预处理
def preprocess_data(stock_list, daily_data, financial_data):
    # 合并数据
    df = pd.merge(daily_data, financial_data, on=['ts_code', 'trade_date'], how='left')
    df = pd.merge(df, stock_list[['ts_code', 'name', 'industry']], on='ts_code', how='left')

    # 转换日期格式
    df['trade_date'] = pd.to_datetime(df['trade_date'])

    # 按股票代码和日期排序
    df = df.sort_values(['ts_code', 'trade_date'])

    # 填充缺失值
    df = df.fillna(method='ffill')

    return df


# 策略1: 传统量价策略 - MACD+RSI组合策略
def strategy_macd_rsi(df, short_window=12, long_window=26, signal_window=9, rsi_period=14,
                      rsi_overbought=70, rsi_oversold=30):
    result = {}

    for stock in df['ts_code'].unique():
        stock_data = df[df['ts_code'] == stock].copy()
        stock_data = stock_data.sort_values('trade_date')

        # 计算移动平均线
        stock_data['ma_short'] = stock_data['close'].rolling(window=5).mean()
        stock_data['ma_long'] = stock_data['close'].rolling(window=20).mean()

        # 计算MACD
        stock_data['ema_short'] = stock_data['close'].ewm(span=short_window, adjust=False).mean()
        stock_data['ema_long'] = stock_data['close'].ewm(span=long_window, adjust=False).mean()
        stock_data['macd'] = stock_data['ema_short'] - stock_data['ema_long']
        stock_data['signal'] = stock_data['macd'].ewm(span=signal_window, adjust=False).mean()
        stock_data['histogram'] = stock_data['macd'] - stock_data['signal']

        # 计算RSI
        delta = stock_data['close'].diff()
        gain = delta.where(delta > 0, 0).rolling(window=rsi_period).mean()
        loss = -delta.where(delta < 0, 0).rolling(window=rsi_period).mean()
        rs = gain / loss
        stock_data['rsi'] = 100 - (100 / (1 + rs))

        # 生成交易信号
        stock_data['macd_signal'] = 0
        stock_data.loc[stock_data['macd'] > stock_data['signal'], 'macd_signal'] = 1
        stock_data.loc[stock_data['macd'] < stock_data['signal'], 'macd_signal'] = -1

        stock_data['rsi_signal'] = 0
        stock_data.loc[stock_data['rsi'] < rsi_oversold, 'rsi_signal'] = 1
        stock_data.loc[stock_data['rsi'] > rsi_overbought, 'rsi_signal'] = -1

        # 综合信号：MACD金叉且RSI超卖时买入，MACD死叉且RSI超买时卖出
        stock_data['signal'] = 0
        stock_data.loc[(stock_data['macd_signal'] == 1) & (stock_data['rsi_signal'] == 1), 'signal'] = 1
        stock_data.loc[(stock_data['macd_signal'] == -1) & (stock_data['rsi_signal'] == -1), 'signal'] = -1

        # 计算持仓
        stock_data['position'] = stock_data['signal'].replace(to_replace=0, method='ffill')
        stock_data['position'] = stock_data['position'].fillna(0)

        # 计算每日收益
        stock_data['returns'] = stock_data['close'].pct_change()
        stock_data['strategy_returns'] = stock_data['position'].shift(1) * stock_data['returns']

        # 计算累计收益
        stock_data['cumulative_returns'] = (1 + stock_data['returns']).cumprod() - 1
        stock_data['strategy_cumulative_returns'] = (1 + stock_data['strategy_returns']).cumprod() - 1

        result[stock] = stock_data

    return result


# 策略2: 基于财务指标的量化策略
def strategy_financial(df, ma_short=5, ma_long=20, pe_threshold=30, pb_threshold=3):
    result = {}

    for stock in df['ts_code'].unique():
        stock_data = df[df['ts_code'] == stock].copy()
        stock_data = stock_data.sort_values('trade_date')

        # 计算移动平均线
        stock_data['ma_short'] = stock_data['close'].rolling(window=ma_short).mean()
        stock_data['ma_long'] = stock_data['close'].rolling(window=ma_long).mean()

        # 生成交易信号
        stock_data['price_signal'] = 0
        stock_data.loc[stock_data['close'] > stock_data['ma_long'], 'price_signal'] = 1
        stock_data.loc[stock_data['close'] < stock_data['ma_short'], 'price_signal'] = -1

        stock_data['value_signal'] = 0
        stock_data.loc[(stock_data['pe'] < pe_threshold) & (stock_data['pb'] < pb_threshold), 'value_signal'] = 1

        # 综合信号：价格突破长期均线且估值合理时买入，价格跌破短期均线时卖出
        stock_data['signal'] = 0
        stock_data.loc[(stock_data['price_signal'] == 1) & (stock_data['value_signal'] == 1), 'signal'] = 1
        stock_data.loc[stock_data['price_signal'] == -1, 'signal'] = -1

        # 计算持仓
        stock_data['position'] = stock_data['signal'].replace(to_replace=0, method='ffill')
        stock_data['position'] = stock_data['position'].fillna(0)

        # 计算每日收益
        stock_data['returns'] = stock_data['close'].pct_change()
        stock_data['strategy_returns'] = stock_data['position'].shift(1) * stock_data['returns']

        # 计算累计收益
        stock_data['cumulative_returns'] = (1 + stock_data['returns']).cumprod() - 1
        stock_data['strategy_cumulative_returns'] = (1 + stock_data['strategy_returns']).cumprod() - 1

        result[stock] = stock_data

    return result


# 策略3: 基于机器学习的量化策略
def strategy_machine_learning(df, lookback=5, test_size=0.2):
    result = {}

    for stock in df['ts_code'].unique():
        stock_data = df[df['ts_code'] == stock].copy()
        stock_data = stock_data.sort_values('trade_date')

        # 计算技术指标作为特征
        stock_data['ma5'] = stock_data['close'].rolling(window=5).mean()
        stock_data['ma10'] = stock_data['close'].rolling(window=10).mean()
        stock_data['ma20'] = stock_data['close'].rolling(window=20).mean()

        # 添加缺失的ma_short和ma_long列
        stock_data['ma_short'] = stock_data['ma5']
        stock_data['ma_long'] = stock_data['ma20']

        # 相对强弱指标
        delta = stock_data['close'].diff()
        gain = delta.where(delta > 0, 0).rolling(window=14).mean()
        loss = -delta.where(delta < 0, 0).rolling(window=14).mean()
        rs = gain / loss
        stock_data['rsi'] = 100 - (100 / (1 + rs))

        # MACD
        stock_data['ema12'] = stock_data['close'].ewm(span=12, adjust=False).mean()
        stock_data['ema26'] = stock_data['close'].ewm(span=26, adjust=False).mean()
        stock_data['macd'] = stock_data['ema12'] - stock_data['ema26']
        stock_data['signal'] = stock_data['macd'].ewm(span=9, adjust=False).mean()
        stock_data['histogram'] = stock_data['macd'] - stock_data['signal']

        # 成交量指标
        stock_data['vol_ma5'] = stock_data['vol'].rolling(window=5).mean()
        stock_data['vol_ma10'] = stock_data['vol'].rolling(window=10).mean()

        # 价格动量
        stock_data['momentum_5'] = stock_data['close'].pct_change(periods=5)
        stock_data['momentum_10'] = stock_data['close'].pct_change(periods=10)

        # 波动率
        stock_data['volatility_5'] = stock_data['close'].rolling(window=5).std()
        stock_data['volatility_10'] = stock_data['close'].rolling(window=10).std()

        # 布林带指标
        stock_data['bb_upper'], stock_data['bb_middle'], stock_data['bb_lower'] = talib.BBANDS(
            stock_data['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)
        stock_data['bb_width'] = (stock_data['bb_upper'] - stock_data['bb_lower']) / stock_data['bb_middle']

        # 成交量变化率
        stock_data['vol_change'] = stock_data['vol'].pct_change()

        # 财务指标
        stock_data['pe'] = stock_data['pe'].fillna(method='ffill')
        stock_data['pb'] = stock_data['pb'].fillna(method='ffill')

        # 创建目标变量：未来n天的收益方向
        stock_data['future_return'] = stock_data['close'].pct_change(lookback).shift(-lookback)
        stock_data['target'] = np.where(stock_data['future_return'] > 0, 1, 0)

        # 丢弃NaN值
        stock_data = stock_data.dropna()

        # 准备特征和目标
        features = [
            'ma5', 'ma10', 'ma20', 'rsi', 'macd', 'signal', 'histogram',
            'vol_ma5', 'vol_ma10', 'momentum_5', 'momentum_10',
            'volatility_5', 'volatility_10', 'bb_width', 'vol_change',
            'pe', 'pb'
        ]

        X = stock_data[features]
        y = stock_data['target']

        # 划分训练集和测试集
        split_index = int(len(stock_data) * (1 - test_size))
        X_train, X_test = X[:split_index], X[split_index:]
        y_train, y_test = y[:split_index], y[split_index:]

        # 训练随机森林模型
        model = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
        model.fit(X_train, y_train)

        # 预测
        stock_data['prediction'] = np.nan
        stock_data.loc[X_test.index, 'prediction'] = model.predict(X_test)

        # 生成交易信号
        stock_data['signal'] = 0
        stock_data.loc[stock_data['prediction'] == 1, 'signal'] = 1
        stock_data.loc[stock_data['prediction'] == 0, 'signal'] = -1

        # 计算持仓
        stock_data['position'] = stock_data['signal'].replace(to_replace=0, method='ffill')
        stock_data['position'] = stock_data['position'].fillna(0)

        # 计算每日收益
        stock_data['returns'] = stock_data['close'].pct_change()
        stock_data['strategy_returns'] = stock_data['position'].shift(1) * stock_data['returns']

        # 计算累计收益
        stock_data['cumulative_returns'] = (1 + stock_data['returns']).cumprod() - 1
        stock_data['strategy_cumulative_returns'] = (1 + stock_data['strategy_returns']).cumprod() - 1

        result[stock] = stock_data

    return result


# 止损止盈功能
def apply_stop_loss_take_profit(result, stop_loss=0.08, take_profit=0.15):
    for stock, data in result.items():
        # 创建新列存储止损止盈后的信号
        data['position_sl_tp'] = data['position'].copy()

        # 记录持仓成本
        data['cost'] = np.nan

        # 初始化持仓状态
        in_position = False
        cost_price = 0

        for i in range(len(data)):
            if data.iloc[i]['signal'] == 1 and not in_position:
                # 买入信号且当前未持仓
                in_position = True
                cost_price = data.iloc[i]['close']
                data.iloc[i, data.columns.get_loc('cost')] = cost_price

            elif in_position:
                # 已持仓，检查是否触发止损或止盈
                current_price = data.iloc[i]['close']
                data.iloc[i, data.columns.get_loc('cost')] = cost_price

                # 止损
                if (current_price - cost_price) / cost_price <= -stop_loss:
                    data.iloc[i, data.columns.get_loc('position_sl_tp')] = 0
                    in_position = False

                # 止盈
                elif (current_price - cost_price) / cost_price >= take_profit:
                    data.iloc[i, data.columns.get_loc('position_sl_tp')] = 0
                    in_position = False

        # 计算考虑止损止盈后的收益
        data['strategy_returns_sl_tp'] = data['position_sl_tp'].shift(1) * data['returns']
        data['strategy_cumulative_returns_sl_tp'] = (1 + data['strategy_returns_sl_tp']).cumprod() - 1

        result[stock] = data

    return result


# 回测和评估
def backtest_and_evaluate(result):
    evaluation = {}

    for stock, data in result.items():
        # 计算评估指标
        total_days = len(data)
        trading_days = len(data[data['position'] != 0])

        # 计算年化收益率
        days_per_year = 252
        total_return = data['strategy_cumulative_returns'].iloc[-1]
        annual_return = (1 + total_return) ** (days_per_year / total_days) - 1

        # 计算最大回撤
        data['cumulative_max'] = data['strategy_cumulative_returns'].cummax()
        data['drawdown'] = data['cumulative_max'] - data['strategy_cumulative_returns']
        max_drawdown = data['drawdown'].max()

        # 计算夏普比率 (假设无风险利率为0)
        risk_free_rate = 0
        sharpe_ratio = (annual_return - risk_free_rate) / (data['strategy_returns'].std() * np.sqrt(days_per_year))

        # 计算胜率
        trades = data[data['signal'] != 0].copy()
        if len(trades) > 0:
            trades['next_return'] = trades['returns'].shift(-1)
            winning_trades = trades[trades['next_return'] > 0]
            win_rate = len(winning_trades) / len(trades)
        else:
            win_rate = 0

        # 存储评估结果
        evaluation[stock] = {
            'total_return': total_return,
            'annual_return': annual_return,
            'max_drawdown': max_drawdown,
            'sharpe_ratio': sharpe_ratio,
            'win_rate': win_rate,
            'trading_days': trading_days,
            'total_days': total_days
        }

    return evaluation


# 可视化函数
def visualize_results(result, evaluation, strategy_name):
    # 1. 绘制累计收益对比图
    plt.figure(figsize=(15, 8))

    for stock, data in result.items():
        plt.plot(data['trade_date'], data['cumulative_returns'], label=f'{stock} - Buy & Hold', linestyle='--')
        plt.plot(data['trade_date'], data['strategy_cumulative_returns'], label=f'{stock} - {strategy_name}')

    plt.title(f'{strategy_name} 策略累计收益对比', fontsize=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('累计收益率', fontsize=12)
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.show()

    # 2. 绘制单只股票的详细分析图
    stock = list(result.keys())[0]
    data = result[stock]

    fig = plt.figure(figsize=(18, 12))
    gs = GridSpec(4, 1, height_ratios=[3, 1, 1, 1])

    # 价格和交易信号
    ax1 = plt.subplot(gs[0])
    ax1.plot(data['trade_date'], data['close'], label='收盘价', color='black', linewidth=1.5)
    ax1.plot(data['trade_date'], data['ma_short'], label=f'{5}日均线', color='blue', linewidth=1)
    ax1.plot(data['trade_date'], data['ma_long'], label=f'{20}日均线', color='red', linewidth=1)

    # 标记买入和卖出信号
    buy_signals = data[data['signal'] == 1]
    sell_signals = data[data['signal'] == -1]

    ax1.scatter(buy_signals['trade_date'], buy_signals['close'], marker='^', color='green', s=100, label='买入信号')
    ax1.scatter(sell_signals['trade_date'], sell_signals['close'], marker='v', color='red', s=100, label='卖出信号')

    ax1.set_title(f'{stock} 价格走势与交易信号', fontsize=15)
    ax1.set_ylabel('价格', fontsize=12)
    ax1.grid(True, alpha=0.3)
    ax1.legend()

    # 持仓情况
    ax2 = plt.subplot(gs[1], sharex=ax1)
    ax2.fill_between(data['trade_date'], 0, data['position'], color='blue', alpha=0.3)
    ax2.set_ylabel('持仓', fontsize=12)
    ax2.grid(True, alpha=0.3)

    # 策略收益与基准收益
    ax3 = plt.subplot(gs[2], sharex=ax1)
    ax3.plot(data['trade_date'], data['cumulative_returns'], label='买入持有收益', color='black')
    ax3.plot(data['trade_date'], data['strategy_cumulative_returns'], label=f'{strategy_name}收益', color='blue')
    ax3.set_ylabel('累计收益', fontsize=12)
    ax3.grid(True, alpha=0.3)
    ax3.legend()

    # 回撤
    ax4 = plt.subplot(gs[3], sharex=ax1)
    ax4.fill_between(data['trade_date'], 0, data['drawdown'], color='red', alpha=0.3)
    ax4.set_ylabel('回撤', fontsize=12)
    ax4.set_xlabel('日期', fontsize=12)
    ax4.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()

    # 3. 使用Plotly创建交互式图表
    fig = make_subplots(
        rows=4, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.05,
        row_heights=[0.4, 0.2, 0.2, 0.2]
    )

    # 价格和移动平均线
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['close'], name='收盘价', line=dict(color='black')),
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['ma_short'], name=f'{5}日均线', line=dict(color='blue')),
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['ma_long'], name=f'{20}日均线', line=dict(color='red')),
        row=1, col=1
    )

    # 买入和卖出信号
    fig.add_trace(
        go.Scatter(
            x=buy_signals['trade_date'],
            y=buy_signals['close'],
            name='买入信号',
            mode='markers',
            marker=dict(size=10, color='green', symbol='triangle-up')
        ),
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(
            x=sell_signals['trade_date'],
            y=sell_signals['close'],
            name='卖出信号',
            mode='markers',
            marker=dict(size=10, color='red', symbol='triangle-down')
        ),
        row=1, col=1
    )

    # 持仓情况
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['position'], name='持仓', fill='tozeroy',
                   fillcolor='rgba(0, 128, 255, 0.3)'),
        row=2, col=1
    )

    # 收益情况
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['cumulative_returns'], name='买入持有收益', line=dict(color='black')),
        row=3, col=1
    )
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['strategy_cumulative_returns'], name=f'{strategy_name}收益',
                   line=dict(color='blue')),
        row=3, col=1
    )

    # 回撤
    fig.add_trace(
        go.Scatter(x=data['trade_date'], y=data['drawdown'], name='回撤', fill='tozeroy',
                   fillcolor='rgba(255, 0, 0, 0.3)'),
        row=4, col=1
    )

    fig.update_layout(
        title=f'{stock} {strategy_name}策略分析',
        height=800,
        xaxis_title='日期',
        yaxis_title='价格',
        yaxis2_title='持仓',
        yaxis3_title='累计收益',
        yaxis4_title='回撤',
        showlegend=True
    )

    fig.show()

    # 4. 策略评估指标雷达图
    eval_df = pd.DataFrame(evaluation).T

    # 标准化指标以便在雷达图中显示
    normalized_eval = pd.DataFrame()
    normalized_eval['总收益率'] = eval_df['total_return'] / eval_df['total_return'].max()
    normalized_eval['年化收益率'] = eval_df['annual_return'] / eval_df['annual_return'].max()
    normalized_eval['夏普比率'] = (eval_df['sharpe_ratio'] - eval_df['sharpe_ratio'].min()) / (
            eval_df['sharpe_ratio'].max() - eval_df['sharpe_ratio'].min())
    normalized_eval['胜率'] = eval_df['win_rate']
    normalized_eval['最大回撤'] = 1 - (eval_df['max_drawdown'] - eval_df['max_drawdown'].min()) / (
            eval_df['max_drawdown'].max() - eval_df['max_drawdown'].min())

    # 为每个股票绘制雷达图
    categories = list(normalized_eval.columns)
    N = len(categories)

    for stock in normalized_eval.index:
        values = normalized_eval.loc[stock].values.flatten().tolist()
        values += values[:1]  # 闭合雷达图

        angles = [n / float(N) * 2 * np.pi for n in range(N)]
        angles += angles[:1]  # 闭合雷达图

        plt.figure(figsize=(8, 8))
        ax = plt.subplot(111, polar=True)
        plt.xticks(angles[:-1], categories, color='grey', size=10)
        ax.set_rlabel_position(0)
        plt.yticks([0.2, 0.4, 0.6, 0.8, 1], ["0.2", "0.4", "0.6", "0.8", "1.0"], color="grey", size=7)
        plt.ylim(0, 1)

        ax.plot(angles, values, linewidth=2, linestyle='solid')
        ax.fill(angles, values, 'b', alpha=0.1)

        plt.title(f'{stock} {strategy_name}策略评估雷达图', size=15, y=1.1)
        plt.show()

    # 5. 热力图比较不同股票的策略表现
    plt.figure(figsize=(12, 8))
    heatmap_data = eval_df[['total_return', 'annual_return', 'max_drawdown', 'sharpe_ratio', 'win_rate']]
    heatmap_data.columns = ['总收益率', '年化收益率', '最大回撤', '夏普比率', '胜率']

    # 反转最大回撤的值，使其越高越好
    heatmap_data['最大回撤'] = -heatmap_data['最大回撤']

    # 标准化数据
    heatmap_data = (heatmap_data - heatmap_data.mean()) / heatmap_data.std()

    # 绘制热力图
    sns.heatmap(heatmap_data, annot=True, cmap='RdYlGn', fmt='.2f', linewidths=.5)
    plt.title(f'{strategy_name}策略各股票表现热力图', fontsize=15)
    plt.tight_layout()
    plt.show()

    # 6. 各策略综合比较条形图
    plt.figure(figsize=(15, 10))

    metrics = ['total_return', 'annual_return', 'max_drawdown', 'sharpe_ratio', 'win_rate']
    metric_names = ['总收益率', '年化收益率', '最大回撤', '夏普比率', '胜率']

    for i, metric in enumerate(metrics):
        plt.subplot(2, 3, i + 1)
        values = [eval[stock][metric] for eval in [evaluation] for stock in eval.keys()]
        labels = [f"{stock}" for eval in [evaluation] for stock in eval.keys()]

        if metric == 'max_drawdown':
            # 最大回撤越小越好，所以我们用负值显示
            values = [-v for v in values]
            plt.bar(labels, values, color='red', alpha=0.6)
            plt.title(f'{metric_names[i]}')
        else:
            plt.bar(labels, values, color='blue', alpha=0.6)
            plt.title(f'{metric_names[i]}')

        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()


# 新增策略对比可视化函数
def compare_strategies_visualization(results, evaluations, strategy_names):
    # 1. 多策略累计收益对比图（所有股票）
    plt.figure(figsize=(18, 10))

    for i, strategy in enumerate(strategy_names):
        for stock in results[strategy].keys():
            data = results[strategy][stock]
            plt.plot(data['trade_date'], data['strategy_cumulative_returns'],
                     label=f'{stock} - {strategy}', alpha=0.7, linewidth=1.5)

    plt.title('多策略累计收益对比', fontsize=18)
    plt.xlabel('日期', fontsize=14)
    plt.ylabel('累计收益率', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.legend(ncol=2, fontsize=12, loc='upper left')
    plt.tight_layout()
    plt.show()

    # 2. 单只股票多策略对比图（以第一只股票为例）
    stock = list(results[strategy_names[0]].keys())[0]
    fig = plt.figure(figsize=(18, 15))
    gs = GridSpec(5, 1, height_ratios=[3, 1, 1, 1, 2])

    # 价格与策略信号
    ax1 = plt.subplot(gs[0])
    data = results[strategy_names[0]][stock]
    ax1.plot(data['trade_date'], data['close'], label='收盘价', color='black', linewidth=2)
    ax1.plot(data['trade_date'], data['ma_short'], label=f'{5}日均线', color='blue', linewidth=1)
    ax1.plot(data['trade_date'], data['ma_long'], label=f'{20}日均线', color='red', linewidth=1)

    for strategy, color in zip(strategy_names, ['blue', 'green', 'red']):
        sig_data = results[strategy][stock]
        buy = sig_data[sig_data['signal'] == 1]
        sell = sig_data[sig_data['signal'] == -1]
        ax1.scatter(buy['trade_date'], buy['close'], marker='^', color=color, s=150,
                    label=f'{strategy} 买入', alpha=0.8)
        ax1.scatter(sell['trade_date'], sell['close'], marker='v', color=color, s=150,
                    label=f'{strategy} 卖出', alpha=0.8)

    ax1.set_title(f'{stock} 多策略交易信号对比', fontsize=16)
    ax1.set_ylabel('价格', fontsize=14)
    ax1.grid(True, alpha=0.3)
    ax1.legend(ncol=3, fontsize=12, loc='upper left')

    # 持仓对比
    ax2 = plt.subplot(gs[1], sharex=ax1)
    for strategy, color in zip(strategy_names, ['blue', 'green', 'red']):
        pos_data = results[strategy][stock]
        ax2.plot(pos_data['trade_date'], pos_data['position'], label=strategy,
                 color=color, linewidth=2, alpha=0.7)

    ax2.set_ylabel('持仓', fontsize=14)
    ax2.grid(True, alpha=0.3)
    ax2.legend()

    # 收益曲线对比
    ax3 = plt.subplot(gs[2], sharex=ax1)
    for strategy, color in zip(strategy_names, ['blue', 'green', 'red']):
        ret_data = results[strategy][stock]
        ax3.plot(ret_data['trade_date'], ret_data['strategy_cumulative_returns'],
                 label=strategy, color=color, linewidth=2, alpha=0.9)

    ax3.set_ylabel('累计收益', fontsize=14)
    ax3.grid(True, alpha=0.3)
    ax3.legend()

    # 回撤对比
    ax4 = plt.subplot(gs[3], sharex=ax1)
    for strategy, color in zip(strategy_names, ['blue', 'green', 'red']):
        dd_data = results[strategy][stock]
        ax4.plot(dd_data['trade_date'], dd_data['drawdown'],
                 label=strategy, color=color, linewidth=2, alpha=0.7)

    ax4.set_ylabel('回撤', fontsize=14)
    ax4.set_xlabel('日期', fontsize=14)
    ax4.grid(True, alpha=0.3)
    ax4.legend()

    # 指标雷达图对比
    ax5 = plt.subplot(gs[4], polar=True)
    metrics = ['总收益率', '年化收益率', '夏普比率', '胜率', '最大回撤']
    N = len(metrics)
    angles = np.linspace(0, 2 * np.pi, N, endpoint=False).tolist()
    angles += angles[:1]

    for i, strategy in enumerate(strategy_names):
        eval_data = list(evaluations[i].values())[0]
        values = [
            eval_data['total_return'],
            eval_data['annual_return'],
            eval_data['sharpe_ratio'],
            eval_data['win_rate'],
            -eval_data['max_drawdown']  # 转换为正值比较
        ]
        values += values[:1]
        ax5.plot(angles, values, linewidth=2, label=strategy, alpha=0.8)
        ax5.fill(angles, values, alpha=0.1)

    ax5.set_xticks(angles[:-1])
    ax5.set_xticklabels(metrics, fontsize=12)
    ax5.set_rlabel_position(45)
    ax5.set_yticks([0.2, 0.4, 0.6, 0.8])
    ax5.set_yticklabels(['0.2', '0.4', '0.6', '0.8'], fontsize=10)
    ax5.set_title('策略关键指标雷达图对比', fontsize=16, y=1.1)
    ax5.legend(loc='upper right', bbox_to_anchor=(1.3, 1))

    plt.tight_layout()
    plt.show()

    # 3. 综合指标柱状图对比
    plt.figure(figsize=(16, 8))
    metrics = ['annual_return', 'sharpe_ratio', 'win_rate', 'max_drawdown']
    metric_labels = ['年化收益率', '夏普比率', '胜率', '最大回撤']

    for i, metric in enumerate(metrics):
        plt.subplot(2, 2, i + 1)
        for j, strategy in enumerate(strategy_names):
            data = [eval[stock][metric] for stock, eval in evaluations[j].items()]
            if metric == 'max_drawdown':
                data = [-d for d in data]  # 最大回撤取反
            x = np.arange(len(data)) + j * 0.25
            plt.bar(x, data, width=0.2, label=strategy_names[j])

        plt.title(metric_labels[i], fontsize=14)
        plt.ylabel('值', fontsize=12)
        plt.xticks(np.arange(len(data)) + 0.375, [stock.split('.')[0] for stock in evaluations[0].keys()])
        plt.legend()
        plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()


# 修改主函数中的策略比较部分
def run_and_compare_strategies(df):
    # 运行三个策略
    result_strategy1 = strategy_macd_rsi(df)
    result_strategy1 = apply_stop_loss_take_profit(result_strategy1)
    evaluation_strategy1 = backtest_and_evaluate(result_strategy1)

    result_strategy2 = strategy_financial(df)
    result_strategy2 = apply_stop_loss_take_profit(result_strategy2)
    evaluation_strategy2 = backtest_and_evaluate(result_strategy2)

    result_strategy3 = strategy_machine_learning(df)
    result_strategy3 = apply_stop_loss_take_profit(result_strategy3)
    evaluation_strategy3 = backtest_and_evaluate(result_strategy3)

    # 组合结果和评估数据
    results = {
        'MACD+RSI': result_strategy1,
        '财务指标': result_strategy2,
        '机器学习': result_strategy3
    }
    evaluations = [
        evaluation_strategy1,
        evaluation_strategy2,
        evaluation_strategy3
    ]
    strategy_names = list(results.keys())

    # 绘制原策略可视化
    for strategy in strategy_names:
        visualize_results(results[strategy], evaluations[strategy_names.index(strategy)], strategy)

    # 新增多策略对比可视化
    compare_strategies_visualization(results, evaluations, strategy_names)

    # 原策略比较图
    plt.figure(figsize=(15, 5))

    strategies = ['MACD+RSI', '财务指标', '机器学习']
    evaluations = [evaluation_strategy1, evaluation_strategy2, evaluation_strategy3]

    # 计算每种策略的平均表现
    avg_returns = []
    avg_sharpe = []
    avg_max_dd = []

    for eval in evaluations:
        returns = [e['annual_return'] for e in eval.values()]
        sharpe = [e['sharpe_ratio'] for e in eval.values()]
        max_dd = [e['max_drawdown'] for e in eval.values()]

        avg_returns.append(np.mean(returns))
        avg_sharpe.append(np.mean(sharpe))
        avg_max_dd.append(np.mean(max_dd))

    # 绘制策略比较图
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.bar(strategies, avg_returns, color=['blue', 'green', 'red'], alpha=0.6)
    plt.title('平均年化收益率比较')
    plt.ylabel('年化收益率')
    plt.grid(True, alpha=0.3)

    plt.subplot(1, 3, 2)
    plt.bar(strategies, avg_sharpe, color=['blue', 'green', 'red'], alpha=0.6)
    plt.title('平均夏普比率比较')
    plt.ylabel('夏普比率')
    plt.grid(True, alpha=0.3)

    plt.subplot(1, 3, 3)
    plt.bar(strategies, [-dd for dd in avg_max_dd], color=['blue', 'green', 'red'], alpha=0.6)
    plt.title('平均最大回撤比较')
    plt.ylabel('最大回撤(负值表示越小越好)')
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()

    # 返回最佳策略
    best_idx = np.argmax(avg_returns)
    print(f"根据年化收益率，最佳策略是: {strategies[best_idx]}")

    return results


# 主函数
def main():
    # 获取股票列表
    stock_list = get_stock_list()
    print(f"获取到 {len(stock_list)} 只股票")

    # 获取股票数据
    daily_data = get_stock_data(stock_list)
    print(f"获取到 {len(daily_data)} 条日线数据")

    # 获取财务数据
    financial_data = get_financial_data(stock_list)
    print(f"获取到 {len(financial_data)} 条财务数据")

    # 数据预处理
    df = preprocess_data(stock_list, daily_data, financial_data)
    print(f"预处理后的数据有 {len(df)} 条记录")

    # 运行并比较三种策略
    results = run_and_compare_strategies(df)

    # 可以在这里添加更多分析或交易系统优化

    print("量化交易系统运行完毕!")


if __name__ == "__main__":
    main()
